package javassist;

import java.io.IOException;
import java.lang.instrument.ClassFileTransformer;
import java.lang.instrument.IllegalClassFormatException;
import java.lang.instrument.Instrumentation;
import java.lang.reflect.Method;
import java.security.ProtectionDomain;
import javassist.*;
import annotationI.Raaj;

/*
	Author: Gary Benattar
	Email: g.benattar@gmail.com
 */

public class Inspector implements ClassFileTransformer{

	/***
	 * This method is from the interface ClassFileTransformer
	 */
	public byte[] transform(ClassLoader loader, String className, Class<?> clazz,
			ProtectionDomain domain, byte[] bytes)
	throws IllegalClassFormatException {
		String[] ignore = new String[] { "sun/", "java/", "javax/" };
		for (int i = 0; i < ignore.length; i++) {
			if (className.startsWith(ignore[i])) {
				return bytes;
			}
		}
		return inspectClass(className, clazz, bytes);
	}

	/**
	 * This function is called by javassist, here we can add some Timer...
	 * @param agentArgument
	 * @param instrumentation
	 */
	public static void premain(String agentArgument,
			Instrumentation instrumentation) {
		// Add the transform (should implements ClassFileTransformer)
		// from the doc ->
		// When a JVM is launched in a way that indicates an agent class. 
		// In that case an Instrumentation instance is passed to the premain 
		// method of the agent class
		System.out.println("in premain");
		instrumentation.addTransformer(new Inspector());
	}

	/***
	 * Class inspector
	 * @param name
	 * @param clazz
	 * @param b
	 * @return
	 * @throws NotFoundException 
	 * @throws IOException 
	 * @throws ClassNotFoundException 
	 * @throws CannotCompileException 
	 */
	private byte[] inspectClass(String name, Class<?> clazz, byte[] b)
	{
		ClassPool pool = ClassPool.getDefault();
		
		CtClass cl;
		try {
			cl = pool.makeClass(new java.io.ByteArrayInputStream(b));
			
			if (cl.isInterface() == false) {
				CtBehavior[] methods = cl.getDeclaredBehaviors();
				for (int i = 0; i < methods.length; i++) {
					if (!methods[i].isEmpty()) {
						inspectMethod(cl, methods[i]);
					}
				}
			}
			return cl.toBytecode();
		}catch (RuntimeException e) {
			// TODO Auto-generated catch block
			e.printStackTrace();
		} catch (NotFoundException e) {
			// TODO Auto-generated catch block
			e.printStackTrace();
		} catch (CannotCompileException e) {
			// TODO Auto-generated catch block
			e.printStackTrace();
		} catch (ClassNotFoundException e) {
			// TODO Auto-generated catch block
			e.printStackTrace();
		} catch (IOException e) {
			// TODO Auto-generated catch block
			e.printStackTrace();
		}
		return b;
	}

	/***
	 * Method inspector
	 * @param method
	 * @throws NotFoundException
	 * @throws CannotCompileException
	 * @throws ClassNotFoundException 
	 * @throws IOException 
	 */
	private void inspectMethod(CtClass clazz, CtBehavior method) throws NotFoundException
	, CannotCompileException
	, ClassNotFoundException
	, IOException {
		if(!method.isEmpty()){
			Object[] annotations = method.getAnnotations();
			for(Object o: annotations){
				if(o instanceof Raaj){
					if(((Raaj)o).name().equals("safer")){
						int param = ((Raaj)o).param();
						method.insertBefore("{ if($" + param + " == null) return; }");
					}
				}
			}
			clazz.writeFile();
		}
	}
}
